# 采用环差法预测GDP
import pandas as pd
import numpy as np

from data_sets import prepare_xy_data

from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge, Lasso
from sklearn.metrics import mean_absolute_error

from pyecharts import options as opts
from pyecharts.charts import Line


# 设置legend和text格式
tjd_legend_opts = opts.LegendOpts(
    pos_top="7%", 
    is_show=True, 
    textstyle_opts=opts.TextStyleOpts(font_size=18)
    )
tjd_text_opts = opts.TextStyleOpts(font_size=18)


def pred_and_plot(y_name):
    # 准备y和x数据
    y_data, x_data_set = prepare_xy_data(y_name)
    # 将x数据合并为一个DataFrame
    x_data_df = pd.concat([v for k,v in x_data_set.items()], axis=1)

    # 复制y数据
    y_raw = y_data.copy(deep=True)
    # y_raw.to_excel('y_raw.xlsx')
    print(y_data.shape)
    y_data = y_data.diff(1).iloc[1:] # y数据求差分
    x_data_df = x_data_df.diff(1).iloc[1:] # x数据求差分

    # （开始日期和结束日期）
    spilt_date_end = y_data.index[-1]
    spilt_date_start = max(x_data_df.index[0], y_data.index[0])

    X = x_data_df.loc[spilt_date_start:,:]
    print('X shape',X.shape)
    # X.to_excel('X.xlsx')
    y = y_data.loc[spilt_date_start:,:].iloc[:,0]
    print('y shape', y.shape)
    # y.to_excel('y.xlsx')
    X = X.fillna(method='ffill').fillna(method='bfill')
    # 数据标准化

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 训练验证数据切分点
    split_point = int(len(X) * 0.87)
    X_train = X_scaled[:split_point]
    X_test = X_scaled[split_point:len(y)]
    X_test2 = X_scaled[split_point:] # test2数据比test数据多1-2期
    y_train = y[:split_point]
    y_test = y[split_point:]

    # 岭回归模型
    ridge = Ridge()
    # 网格搜索岭参数
    param_grid = {'alpha': np.logspace(-8, 8, 200)}
    grid = GridSearchCV(ridge, param_grid, cv=5, scoring='neg_mean_squared_error')
    grid.fit(X_train, y_train)

    # 最佳模型
    best_ridge = grid.best_estimator_
    # 预测
    y_train_p = best_ridge.predict(X_train)
    y_pred = best_ridge.predict(X_test)
    y_pred2 = best_ridge.predict(X_test2)
    # 评估模型
    mse = mean_absolute_error(y_test, y_pred)
    print(f'Mean Squared Error: {mse}')

    # 解释系数
    coefficients = pd.DataFrame(best_ridge.coef_, X.columns, columns=['Coefficient'])
    y_t = pd.Series(y_train_p, index=X.index[:split_point], name='训练期预测')
    y_p2 = pd.Series(y_pred2, index=X.index[split_point:], name='验证期预测值')

    valid_y = pd.concat([y_test, y_p2], axis=1)
    valid_y.columns = ['环差', '环差P']
    train_y = pd.concat([y_train, y_t], axis=1)
    train_y.columns = ['环差', '环差P']

    train_valid = pd.concat([train_y, valid_y], axis=0)

    df = pd.concat([y_raw, train_valid], axis=1)
    df['预测'] = np.nan
    df['误差'] = np.nan
    # 预测结果等于上期预测值 + 预测的环差
    for idx in range(1, len(train_valid)):
        t = train_valid.index[idx]
        t_last_q = train_valid.index[idx-1]
        df.loc[t, '预测'] = df.loc[t_last_q, y_raw.columns[0]] + df.loc[t, '环差P']
        df.loc[t, '误差'] = abs(df.loc[t, '预测'] - df.loc[t, y_raw.columns[0]])
    df = df.round(1)

    # 画图
    b = Line()
    b.add_xaxis([s.strftime("%Y-%m-%d") for s in df.index]) # x数据（常是索引）
    for y_name in df.columns:
        se = df[y_name]
        b.add_yaxis(y_name, y_axis=se.to_list()) # y数据
        
        b.set_global_opts(
            datazoom_opts=[
                opts.DataZoomOpts(range_start=30, range_end=100),
                opts.DataZoomOpts(type_="inside")],
            toolbox_opts=opts.ToolboxOpts(),
            yaxis_opts=opts.AxisOpts(name="%", axislabel_opts=opts.LabelOpts(font_size=16)),
            xaxis_opts=opts.AxisOpts(name="日期", axislabel_opts=opts.LabelOpts(font_size=16)),
            legend_opts=tjd_legend_opts
            )
    b.render(f'分项加总\预测{y_raw.columns[0]}2025年3月.html')
    pth = f'分项加总\预测{y_raw.columns[0]}2025年3月.xlsx'
    df.to_excel(pth)
    return pth, df

# 要在pred.yaml里面注册过
y_names = [
    'GDP不变价当季同比',
    # 'GDP工业不变价当季同比',
    # 'GDP交通运输不变价当季同比',
    # 'GDP住宿餐饮业不变价当季同比',
    # 'GDP房地产不变价当季同比',
    # 'GDP批发零售不变价当季同比',
    # 'GDP金融不变价当季同比',
    # 'GDP信息传输不变价当季同比',
    # 'GDP商务服务不变价当季同比',
    # 'GDP建筑业不变价当季同比'
    ]
for y_name in y_names:
    pth, df = pred_and_plot(y_name)
    print(pth)